#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ECCSide-Channel Attack - True实Elliptic Curve点坐标泄漏模型

核心实现：
- Attack时重新CalculateElliptic Curve点
- FromFalse设点坐标Extract泄漏（与Power仿True完全Match）
- 理论完全正确的ECCSide-Channel Attack

Calculate优化：
- 只Calculate前几步的点运算（而非Completek*P）
- 缓存Elliptic Curve点CalculateResult
- Support逐字节独立Attack

技术特点：
- Calculate开销大：EachFalse设RequireElliptic Curve点运算
- 泄漏模型与Power仿True完全一致
- Support多种标准曲线
"""

import numpy as np
from typing import Dict, List, Tuple, Optional, Any
import sys
import os

# 本地导入
from simplified_ecc_math import SimplifiedECCPoint, SimplifiedECCCurve, get_curve

# Try相对导入，Failed则Use绝对导入
try:
    from ..base_algorithm_attack import AlgorithmSpecificAttack
    from ...base_side_channel_attack import AttackResult
except (ImportError, ValueError):
    # 相对导入Failed，Use绝对导入
    from algorithm_specific_attacks.base_algorithm_attack import AlgorithmSpecificAttack
    from side_channel_attacks.base_side_channel_attack import AttackResult


class ECCSpecificAttacks(AlgorithmSpecificAttack):
    """
    ECC CPAAttack - True实点坐标泄漏模型
    
    核心技术：
    - 重新CalculateElliptic Curve点，From点坐标Extract泄漏
    - 泄漏模型与Power仿True完全Match
    - Support逐字节独立Attack
    
    泄漏模型必须与Power仿True完全Match
    """
    
    def __init__(self, curve_name: str = "secp128r1"):
        """
        InitializeECCAttack
        
        Args:
            curve_name: Elliptic Curve名称
        """
        super().__init__("ECC")
        
        self.curve_name = curve_name
        self.curve = get_curve(curve_name)
        
        # 新架构：分离的模型（导入ECC模型）
        try:
            from .ecc_models import ECCLeakageModels
        except (ImportError, ValueError):
            from ecc_models import ECCLeakageModels
        self.intermediate_functions = ECCLeakageModels.get_intermediate_functions()
        self.leakage_models = ECCLeakageModels.get_leakage_models()
        self.operation_mapping = ECCLeakageModels.get_operation_to_intermediate_mapping()
        
        # Data
        self.traces = None
        self.plaintexts = None
        self.keys = None
        
        # 多字节AttackStatus
        self.recovered_key_bytes = {}
        self.recovery_confidence = {}
        
        # 缓存：Avoid重复Calculate
        self.point_cache = {}  # (message_scalar, hypothesis) -> Point
        
        # 优化Options
        self.use_cache = True
        self.attack_depth = 1  # Attack深度：1=Step 1，2=前两步，etc.
    
    def get_intermediate_functions(self) -> Dict[str, Any]:
        """获取Intermediate valueCalculateFunction"""
        return self.intermediate_functions
    
    def get_leakage_models(self) -> Dict[str, Any]:
        """获取泄漏表示模型"""
        return self.leakage_models
    
    def create_leakage_models(self) -> Dict[str, Any]:
        """CreateECC特定的泄漏模型（用于基Class兼容）"""
        combined = {}
        combined.update(self.intermediate_functions)
        combined.update(self.leakage_models)
        return combined
    
    def get_operation_to_intermediate_mapping(self) -> Dict[str, str]:
        """获取目标操作ToIntermediate valueFunction的Mapping"""
        return self.operation_mapping
    
    def multi_byte_joint_attack(self, target_bytes: List[int], 
                                target_operation: str = 'scalar_multiplication') -> Dict[int, int]:
        """
        多字节联合Attack - V3Version
        
        Note：V3由于Calculate开销大，建议逐字节AttackOrUse优化策略
        
        Args:
            target_bytes: 目标字节位置List
            target_operation: 目标操作
        
        Returns:
            recovered_bytes: {字节位置: 恢复的值}
        """
        recovered_bytes = {}
        known_bytes = {}
        
        print(f"\nStartV3多字节Attack（True实点坐标泄漏）")
        print(f"目标字节: {target_bytes}")
        print(f"曲线: {self.curve}")
        
        # 逐字节Attack（V3推荐方式）
        for byte_pos in target_bytes:
            print(f"\nAttack字节 {byte_pos}...")
            
            result = self._attack_single_byte_v3(byte_pos, target_operation, known_bytes)
            
            if result.success:
                recovered_bytes[byte_pos] = result.best_hypothesis
                known_bytes[byte_pos] = result.best_hypothesis
                self.recovery_confidence[byte_pos] = result.confidence
                
                print(f"  [OK] 恢复Success: 0x{result.best_hypothesis:02x}, 置信度: {result.confidence:.4f}")
            else:
                print(f"  [FAIL] 恢复Failed")
        
        return recovered_bytes
    
    def _attack_single_byte_v3(self, target_byte: int, target_operation: str,
                                known_bytes: Dict[int, int]) -> AttackResult:
        """
        单字节Attack - V3Version（重新CalculateElliptic Curve点）
        
        核心流程：
        1. 对EachKeyFalse设（0-255）
        2. 重新CalculateElliptic Curve点操作
        3. FromFalse设点坐标Extract泄漏
        4. Calculate与PowerTrace的相关性
        5. 选择相关性最高的False设
        
        Args:
            target_byte: 目标字节位置
            target_operation: 目标操作
            known_bytes: 已知的Key字节
        
        Returns:
            AttackResult
        """
        # 手动实现CPAAttack（因ForRequire自定义泄漏模型）
        best_hypothesis = 0
        best_correlation = -1.0
        correlations = []
        
        # 对EachFalse设（0-255）
        for hypothesis in range(256):
            # CalculateIntermediate value（V3：FromTrue实点坐标Extract）
            intermediate_values = self._calculate_intermediate_values_v3(
                self.plaintexts[:, target_byte],
                hypothesis,
                target_byte,
                target_operation,
                known_bytes
            )
            
            # Calculate与PowerTrace的平均相关性
            avg_correlation = 0.0
            valid_points = 0
            
            # Check泄漏值YesNoHas变化
            if np.std(intermediate_values) < 1e-10:
                # 泄漏值没Has变化，Skip这个False设
                correlations.append(0.0)
                continue
            
            # 对Trace的Each时间点Calculate相关性
            for time_point in range(self.traces.shape[1]):
                power_samples = self.traces[:, time_point]
                
                # CheckPower样本YesNoHas变化
                if np.std(power_samples) < 1e-10:
                    continue
                
                # Calculate皮尔逊相关系数
                correlation = np.corrcoef(intermediate_values, power_samples)[0, 1]
                
                if not np.isnan(correlation):
                    avg_correlation += abs(correlation)
                    valid_points += 1
            
            if valid_points > 0:
                avg_correlation /= valid_points
                correlations.append(avg_correlation)
                
                if avg_correlation > best_correlation:
                    best_correlation = avg_correlation
                    best_hypothesis = hypothesis
            else:
                correlations.append(0.0)
        
        # Calculate置信度（标准Method：最大值与次大值的比值）
        correlations = np.array(correlations)
        if correlations.max() > 0:
            sorted_corr = np.sort(correlations)[::-1]  # 降序排列
            if len(sorted_corr) > 1 and sorted_corr[1] > 1e-10:
                # 标准Method：最佳相关性 / 次佳相关性
                confidence = sorted_corr[0] / sorted_corr[1]
                confidence = min(confidence / 2.0, 1.0)  # NormalizationTo[0,1]
            else:
                # If只Has一个明显的候选，Use绝对值
                confidence = min(abs(best_correlation), 1.0)
        else:
            confidence = 0.0
        
        # ReturnResult
        result = AttackResult(
            attack_type='CPA',
            target_byte=target_byte,
            execution_time=0.0,  # 稍后会Update
            success=(confidence > 0.1),
            best_hypothesis=best_hypothesis,
            confidence=confidence,
            metadata={
                'correlation_scores': correlations.tolist(),
                'leakage_model': 'real_point_coordinates_v3'
            }
        )
        
        return result
    
    def _calculate_intermediate_values_v3(self, target_data: np.ndarray, hypothesis: int,
                                          target_byte: int, target_operation: str,
                                          known_bytes: Dict[int, int]) -> np.ndarray:
        """
        CalculateIntermediate value - V3Version（FromTrue实Elliptic Curve点坐标Extract泄漏）
        
        关键Step：
        1. According tomessageAndFalse设Key，CalculateElliptic Curve点
        2. From点坐标Extract泄漏（与V3Power仿TrueMatch）
        3. Return泄漏值（用于相关性Calculate）
        
        Args:
            target_data: 目标Data（message字节）
            hypothesis: False设的Key字节值
            target_byte: CurrentAttack的字节位置
            target_operation: 目标操作
            known_bytes: 已知的Key字节
        
        Returns:
            Intermediate valueArray（泄漏值）
        """
        results = []
        
        for msg_byte in target_data:
            # CalculateElliptic Curve点并Extract泄漏
            leakage = self._compute_point_leakage(
                int(msg_byte), hypothesis, target_byte, known_bytes
            )
            results.append(leakage)
        
        return np.array(results, dtype=np.float64)
    
    def _compute_point_leakage(self, message_byte: int, key_byte: int, 
                               byte_pos: int, known_bytes: Dict[int, int]) -> float:
        """
        CalculateElliptic Curve点泄漏 - V3核心Function（按字节独立修复版）
        
        修复策略（与V3Power仿True完全一致）：
        - Each字节独立Calculate泄漏
        - 泄漏Based on：message_byte * key_byte的PartialScalar乘法
        - Not依赖其他字节，可以逐字节Match
        
        Args:
            message_byte: message的Current字节
            key_byte: False设的Key字节
            byte_pos: 字节位置
            known_bytes: 已知的Key字节（本Method中Not再Require）
        
        Returns:
            泄漏值（平均Power因子）
        """
        # Frommessage_byteGenerateInput点P（与Power仿True完全一致）
        message_bytes = bytes([message_byte] * self.curve.byte_length)
        message_int = int.from_bytes(message_bytes, 'big')
        message_scalar = (message_int % (self.curve.n - 1)) + 1
        P = self.curve.G.scalar_multiply(message_scalar)
        
        # 构建PartialKey（只HasCurrent字节，与Power仿True完全一致）
        partial_key = bytearray(self.curve.byte_length)
        partial_key[byte_pos] = key_byte
        k_int = int.from_bytes(bytes(partial_key), 'big')
        k = (k_int % (self.curve.n - 1)) + 1
        
        # Calculate该字节的独立泄漏（与Power仿True完全一致）
        leakage = self._compute_byte_independent_leakage(P, k)
        
        return leakage
    
    def _compute_byte_independent_leakage(self, P: SimplifiedECCPoint, k: int) -> float:
        """
        CalculateSingle字节的独立泄漏（与V3Power仿True_compute_byte_leakage完全一致）
        
        Args:
            P: Input点
            k: PartialScalar（只HasCurrent字节非0）
        
        Returns:
            平均泄漏值（Power因子）
        """
        k_bits = bin(k)[2:]
        if len(k_bits) < 2:
            return 1.0  # DefaultPower因子
        
        Q = P
        total_leakage = 0.0
        leakage_count = 0
        
        # Calculate前8-10步的泄漏（与Power仿True一致）
        for bit_idx in range(1, min(len(k_bits), 10)):
            # 点倍乘
            Q_old = Q
            Q = Q.double()
            double_leakage = self._extract_double_leakage(Q_old, Q)
            total_leakage += double_leakage
            leakage_count += 1
            
            # 点加法
            if k_bits[bit_idx] == '1':
                Q_old = Q
                Q = Q.add(P)
                add_leakage = self._extract_add_leakage(Q_old, P, Q)
                total_leakage += add_leakage
                leakage_count += 1
        
        return total_leakage / leakage_count if leakage_count > 0 else 1.0
    
    def _compute_full_scalar_mult_leakage(self, P: SimplifiedECCPoint, k: int, 
                                           byte_pos: int) -> float:
        """
        CalculateCompleteScalar乘法k*P，ExtractCurrent字节对应比特位的泄漏
        
        修复：CalculateComplete的Double-and-Add，与V3Power仿True完全一致
        
        Args:
            P: Input点
            k: CompleteScalar（已知字节 + CurrentFalse设 + 其他For0）
            byte_pos: CurrentAttack的字节位置
        
        Returns:
            Current字节对应比特位的平均泄漏
        """
        # ExecuteComplete的Double-and-Add，收集All泄漏
        # Not再限制特定字节，因ForNot同False设会产生Not同的CompleteTrace
        k_bits = bin(k)[2:]
        
        if len(k_bits) < 2:
            return 1.0  # Defaulthw_factor
        
        # ExecuteComplete的Double-and-Add（与V3Power仿True一致）
        Q = P
        total_leakage = 0.0
        leakage_count = 0
        
        # 遍历k的All位（From第2位Start）
        # 收集All操作的泄漏（Not限制特定字节）
        for bit_idx in range(1, min(len(k_bits), 32)):  # 限制前32位Avoid过慢
            # 点倍乘（总YesExecute）
            Q_old = Q
            Q = Q.double()
            
            # 收集泄漏
            double_leakage = self._extract_double_leakage(Q_old, Q)
            total_leakage += double_leakage
            leakage_count += 1
            
            # 点加法（IfCurrent位Yes1）
            if k_bits[bit_idx] == '1':
                Q_old = Q
                Q = Q.add(P)
                
                # 收集泄漏
                add_leakage = self._extract_add_leakage(Q_old, P, Q)
                total_leakage += add_leakage
                leakage_count += 1
        
        # Return平均泄漏
        return total_leakage / leakage_count if leakage_count > 0 else 1.0
    
    def _extract_double_leakage(self, Q: SimplifiedECCPoint, 
                                 result: SimplifiedECCPoint) -> float:
        """
        From点倍乘Extract泄漏 - 必须与V3Power仿True完全Match
        
        泄漏模型（与ecc_power_v3.py完全一致）：
        1. Input点坐标的汉明权重
        2. Output点坐标的汉明权重
        3. XOR差分的汉明权重
        4. 应用hw_factorConvert公式（关键！）
        """
        if Q.is_infinity or result.is_infinity:
            # DefaultPower因子
            return 1.0
        
        q_x_low = Q.x & 0xFF
        q_y_low = Q.y & 0xFF
        result_x_low = result.x & 0xFF
        result_y_low = result.y & 0xFF
        
        # 与V3Power仿True完全相同的泄漏模型
        hw_input = bin(q_x_low ^ q_y_low).count('1')
        hw_output = bin(result_x_low ^ result_y_low).count('1')
        hw_diff = bin(q_x_low ^ result_x_low).count('1')
        
        # 加权组合（与Power仿True一致）
        combined_hw = (hw_input * 0.3 + hw_output * 0.3 + hw_diff * 0.4)
        
        # 应用hw_factorConvert（与Power仿True一致！）
        hw_factor = 0.8 + (combined_hw / 8.0) * 0.4
        
        return hw_factor
    
    def _extract_add_leakage(self, Q: SimplifiedECCPoint, P: SimplifiedECCPoint,
                             result: SimplifiedECCPoint) -> float:
        """
        From点加法Extract泄漏 - 必须与V3Power仿True完全Match
        
        泄漏模型（与ecc_power_v3.py完全一致）：
        1. 两个Input点x坐标的XOR
        2. InputAndOutput的XOR
        3. All坐标的组合
        4. 应用hw_factorConvert公式（关键！）
        """
        if Q.is_infinity or P.is_infinity or result.is_infinity:
            return 1.0
        
        q_x_low = Q.x & 0xFF
        p_x_low = P.x & 0xFF
        result_x_low = result.x & 0xFF
        
        # 与V3Power仿True完全相同的泄漏模型
        hw_input_diff = bin(q_x_low ^ p_x_low).count('1')
        hw_io_diff = bin(q_x_low ^ result_x_low).count('1')
        hw_combined = bin((q_x_low + p_x_low + result_x_low) & 0xFF).count('1')
        
        # 加权组合（与Power仿True一致）
        combined_hw = (hw_input_diff * 0.3 + hw_io_diff * 0.3 + hw_combined * 0.4)
        
        # 应用hw_factorConvert（与Power仿True一致！）
        hw_factor = 0.8 + (combined_hw / 8.0) * 0.4
        
        return hw_factor
    
    def _generate_point_from_message(self, message: bytes) -> SimplifiedECCPoint:
        """FrommessageGenerateElliptic Curve点"""
        message_int = int.from_bytes(message[:self.curve.byte_length], 'big')
        message_scalar = (message_int % (self.curve.n - 1)) + 1
        return self.curve.G.scalar_multiply(message_scalar)
    
    def load_data(self, npz_file: str):
        """
        LoadPowerTraceData
        
        Args:
            npz_file: npzFilePath
        """
        data = np.load(npz_file, allow_pickle=True)
        self.traces = data['traces']
        self.plaintexts = data['plaintexts']
        self.keys = data['keys']
        
        print(f"[OK] 已LoadTraceData:")
        print(f"   Trace数: {len(self.traces)}")
        print(f"   Trace长度: {len(self.traces[0])}")
        print(f"   Plaintext数: {len(self.plaintexts)}")
    
    def get_algorithm_parameters(self) -> Dict[str, Any]:
        """获取ECCParameters"""
        return {
            'algorithm': 'ECC',
            'type': 'asymmetric',
            'curve': self.curve_name,
            'curve_bits': self.curve.bit_length,
            'version': 'V3',
            'leakage_model': 'real_point_coordinates'
        }
    
    def create_leakage_models(self) -> Dict[str, Any]:
        """CreateECC泄漏模型"""
        return {
            'point_double': 'real_coordinates',
            'point_add': 'real_coordinates',
            'version': 'V3_SchemeB'
        }
    
    def get_target_operations(self) -> List[str]:
        """获取可Attack的目标操作"""
        return ['scalar_multiplication', 'point_double', 'point_add']
    
    def _calculate_intermediate_values(self, plaintexts: np.ndarray, 
                                       key_hypothesis: int, byte_position: int,
                                       operation: str) -> np.ndarray:
        """
        CalculateIntermediate value - 基Class要求的抽象Method
        
        InV3中，这个Method被_calculate_intermediate_values_v3替代
        """
        return self._calculate_intermediate_values_v3(
            plaintexts, key_hypothesis, byte_position, operation, {}
        )
    
    def _get_selection_function(self, plaintexts: np.ndarray,
                                key_hypothesis: int, byte_position: int,
                                operation: str) -> np.ndarray:
        """
        获取选择Function - 基Class要求的抽象Method（用于DPA）
        
        V3主要实现CPA，DPA的选择FunctionBased onIntermediate value的最高位
        """
        intermediate_values = self._calculate_intermediate_values(
            plaintexts, key_hypothesis, byte_position, operation
        )
        # 选择Function：Intermediate value > 阈值
        return (intermediate_values > 4.0).astype(np.int32)

if __name__ == "__main__":
    print("=" * 70)
    print("ECCAttackTest - True实点坐标泄漏模型")
    print("=" * 70)
    
    attack = create_ecc_attack("secp128r1")
    print(f"\n曲线: {attack.curve}")
    print(f"曲线阶: {hex(attack.curve.n)}")
    print(f"\n[OK] ECCAttackModuleInitializeSuccess")
